#include "unpack_process.hh"

using namespace std;

unpack_process::unpack_process(vector<uint32_t> initial_data, int asic_num, string root_filename, bool is_base_data, double LSB, int thresh_limit, string save_path, string file_write) : ch(asic_num), amp(asic_num)
{
  this->initial_data = initial_data;
  this->asic_num = asic_num;
  this->LSB = LSB;
  this->thresh_limit = thresh_limit;
  this->save_path = save_path;
  this->file_write = file_write;
  if (is_base_data) {
    init(is_base_data);
  }
  else {
    init_root(root_filename);
    init(is_base_data);
    write_root();
    complete_adding();
  }
}

unpack_process::~unpack_process()
{
    if (adas_hit_file != nullptr && adas_hit_file->IsOpen()) {
        adas_hit_file->Close();
        delete adas_hit_tree;
        delete adas_hit_file;
    }
}

void unpack_process::init(bool is_base_data)
{
  // 1. extract the frames
  find_start_end();
  cout << "Valid event number: " << adc_data.size() << endl;

  // 2. cache the adc words
  for (int i = 0; i < adc_data.size(); i++) { // loop all event frames
    vector<uint32_t> temp_data;
    vector<uint32_t> temp_idx;

    // extract adc channel words and chip indexes
    for (int j = 3; j < adc_data[i].size() - 1;) { // loop all words of one event frame
      if ((adc_data[i][j] & 4294901760) == 1717960704) {// chip frame header found (high 16-bit: 0x6666)
        int ch_nr = adc_data[i][j] & 255; // ch nr
        int chip_id = (adc_data[i][j] & 65280) / 256; // chip id
        for (int k = 1; k < ch_nr; k++) { // [todo] why skip the first channel?
          temp_idx.push_back(chip_id);
          temp_data.push_back(adc_data[i][j + 1 + k]);
        }
        j = j + ch_nr + 1;
      }
      else {
        cout << cYELLOW << "incorrect chip frame word, skip and continue" << cRESET << endl;
        j++;
      }
    }

    if (temp_idx.size() == 0){ // ignore empty frame and continue to next frame
      cout << cRED << "incorrect fec frame found and skipped! the data idx is " << i << cRESET << endl;
    }
    else { // cache the newly extractd adc data together with corresponding trigger_id
      trig.push_back(adc_data[i][1]);
      data_chip_id.push_back(temp_idx);
      data.push_back(temp_data);
    }
  }

  // 3. extract the channel id and adc counts
  spilt_channel_amp(data, adas_chl, adas_amp, adas_chip, is_base_data);

  cout << cBLUE << "channel amp spilt suceess" << cRESET << endl;
  cout << endl;

  // 4. extract baseline mean and rms for pedestal data
  if (is_base_data) {
    vector<vector<vector<uint32_t>>> base(asic_num);
    vector<vector<double>> means(asic_num);
    vector<vector<double>> means128(asic_num);
    vector<vector<double>> rms(asic_num);
    vector<vector<double>> rms_bin(asic_num);
    vector<vector<double>> rms128(asic_num);
    vector<vector<double>> threshold(asic_num);

    for (int i = 0; i < adas_amp.size(); ++i) {
      for (int j = 0; j < adas_amp[i].size(); j += 128) {
        vector<uint32_t> group;
        double k = j / 128;
        group.insert(group.end(), adas_amp[i].begin() + j, adas_amp[i].begin() + j + 128);
        amp[k].push_back(group);
      }
    }

    for (int i = 0; i < asic_num; i++) {
      clac_base_rms(amp[i], base[i], means[i], rms[i], threshold[i], rms_bin[i], rms128[i], means128[i]);
    }

    // generate threshold for fec configuration
    vector<double> th;
    for (int i = 0; i < threshold.size(); i++) {
      for (int j = 0; j < threshold[i].size(); j++) {
        th.push_back(threshold[i][j]);
      }
    }
    threshold_to_gbt(th);

    draw(means, rms);
    write_base_txt(means128, rms128, threshold);
  }
}

void unpack_process::spilt_channel_amp(vector<vector<uint32_t>> data,
                                       vector<vector<uint32_t>> &ch_id,
                                       vector<vector<uint32_t>> &adc,
                                       vector<vector<uint32_t>> &chip_id,
                                       bool is_base_data) {
  for (int i = 0; i < data.size(); i++) {
    vector<uint32_t> temp_channel;
    vector<uint32_t> temp_amp;
    vector<uint32_t> temp_idx;
    for (int j = 0; j < data[i].size(); j++) {
      int ch = 0;
      int ampp = 0;
      int idxx = 0;
      ampp = data[i][j] & 16777215;
      ch = (data[i][j] & 4278190080) / 16777216;
      idxx = data_chip_id[i][j];
      if (ch <= 128) {
        temp_channel.push_back(ch);
        temp_amp.push_back(ampp);
        temp_idx.push_back(idxx);
      }
    }
    ch_id.push_back(temp_channel);
    adc.push_back(temp_amp);
    chip_id.push_back(temp_idx);
  }
}

void unpack_process::write_root()
{
  if (trig.size() != adas_chip.size()) {
      cout << "error read data" << endl;
  }

  for (int i = 0; i < trig.size(); i++) {
    nhits = adas_chl[i].size();
    triggerid = trig[i];
    for (int j = 0; j < nhits; j++) {
        adas_hit_idx[j] = adas_chip[i][j];
        adas_hit_chl[j] = adas_chl[i][j];
        adas_hit_amp[j] = adas_amp[i][j];
    }
    if (adas_hit_tree != nullptr) {
        adas_hit_tree->Fill();
    }
    else {
        cout << "adas_hit_tree is not initialized!" << std::endl;
    }
  }
}

bool unpack_process::init_root(string root_filename)
{
    string filename = root_filename + ".root";
    adas_hit_file = new TFile(filename.c_str(), "recreate");
    adas_hit_tree = new TTree("adas_initial_data", "ADAS Hit");

    adas_hit_tree->Branch("nhits", &nhits, "nhits/I");
    adas_hit_tree->Branch("triggerid", &triggerid, "triggerid/I");
    adas_hit_tree->Branch("adas_hit_idx", adas_hit_idx, "adas_hit_idx[nhits]/I");
    adas_hit_tree->Branch("adas_hit_chl", adas_hit_chl, "adas_hit_chl[nhits]/I");
    adas_hit_tree->Branch("adas_hit_amp", adas_hit_amp, "adas_hit_amp[nhits]/I");

    cout << cGREEN << filename << cRESET << endl;
    cout << endl;

    return true;
}

void unpack_process::complete_adding()
{
    if (adas_hit_tree != nullptr)
    {
        adas_hit_file->cd();
        adas_hit_tree->Write();
        adas_hit_file->Flush();
        adas_hit_file->Close();
    }
    else
    {
        cout << "adas_hit_tree is not initialized!" << std::endl;
    }
}

/* extract the adc frames */
void unpack_process::find_start_end()
{
  vector<uint32_t> data;
  data = initial_data;
  int N = data.size(); // the remaining words in data buffer

  int header = 44047; // the first two bytes in the frame: 0xac0f
  int ender = 1077438054; // the last word in the frame: 0x4038 6666

  int i = 0; // data buffer iterator (in word, 32-bits)
  while (N > 0) {
    // 1. extract info form header word: the fixed bytes and frame size
    int dec = data[i];
    bitset<32> bits(dec);
    bitset<16> hig16(bits.to_string().substr(0, 16));
    bitset<14> low14(bits.to_string().substr(32 - 14, 14));
    int msb = hig16.to_ulong(); // fixed bytes
    int lsb = low14.to_ulong(); // frame size

    vector<uint32_t> temp;
    if (msb == header) { // 2. check header fixed bytes
      if (((i + 2) < data.size()) && ((i + data[i + 2]) <= data.size()) && (data[i + 2] == (lsb / 4 + 2)) && (data[i + data[i + 2] - 1] == ender)) {// 3. check both the remaining words, frame_length word, byte_length and the trailor word of next frame
        if (data[i + 2] > 802 || data[i + 2] == 0) { // [todo] 802 is the max frame length alowed, but the whole check needed?
          cout << cRED << "lose data ! the data idx is " << i << " " << N << cRESET << endl;
          cout << endl;
          i++;
          N--;
          continue;
        }

        //
        start_index.push_back(i);
        end_index.push_back(i + data[i + 2] - 1);
        temp.insert(temp.end(), data.begin() + i, data.begin() + i + data[i + 2]);
        adc_data.push_back(temp);

        int dec = data[i + 2];
        i = i + dec;
        N = N - dec;
      }
      else if (((i + 2) >= data.size()) || ((i + data[i + 2] - 1) > data.size())) { // break if no complete frame in the remaing words
        cout << cRED << "last fec frame is incomplete and it's ignored." << cRESET << endl;
        break;
      } else if (data[i + data[i + 2] - 1] != ender) {
        cout << cRED << "incorrect fec trailor, skip one word and continue." << cRESET << endl;
        i++;
        N--;
      } else {
        cout << cRED << "incorrect fec frame, skip and continue." << cRESET << endl;
        int dec = data[i + 2];
        i = i + dec;
        N = N - dec;
      }
    }
    else {
      cout << cRED << "incorrect fec header, skip one word and continue." << cRESET
           << endl;
      i++;
      N--;
    }
  }
  cout << cBLUE << "fec data spilt suceess" << cRESET << endl;
  cout << endl;
}

void unpack_process::clac_base_rms(vector<vector<uint32_t>> data, vector<vector<uint32_t>> &base, vector<double> &means, vector<double> &rms, vector<double> &threshold, vector<double> &rms_bin, vector<double> &rms128, vector<double> &means128)
{
  for (int i = 0; i < data.size(); i++) { // event loop
    vector<uint32_t> temp_base;
    for (int j = 0; j < 128; j++) { // channel loop
      temp_base.push_back(data[i][j]);
    }
    base.push_back(temp_base);
  }

  int rows = base.size(); // event nr
  int cols = base[0].size(); // ch nr (i.e. 128)

  for (int i = 0; i < cols; i++) {
    double sum = 0;
    for (int j = 0; j < rows; j++) {
      sum += base[j][i];
    }
    means.push_back(sum / rows);
  }

  vector<double> rms_count; // rms (in counts)
  for (int i = 0; i < cols; i++) {
    double sum = 0;
    for (int j = 0; j < rows; j++) {
      sum += pow(base[j][i] - means[i], 2);
    }
    rms_count.push_back(sqrt(sum / rows));
  }

  for (int i = 0; i < cols; i++) { // rms (in fC)
    rms.push_back(rms_count[i] * LSB);
  }

  means128.resize(128);
  rms128.resize(128);
  for (int i = 0; i < cols; i++) {
    means128[i] = means[i];
    rms128[i] = rms[i];
  }
  for (int i = 0; i < means.size(); i++) { // threshold: in unit of sigma 
    threshold.push_back( thresh_limit * rms_count[i] );
  }
  for (int i = 0; i < cols; i++) {
    rms_bin.push_back(rms_count[i]);
  }

  // swap at index 64 fro the channel 65-128
  // it's needed due to the lack of ch_id info
  int n = 0;
  for (int i = 0; i < means.size(); i++) {
    if (i > 63 && n <= 32) {
      int t = 0;
      double r = 0;

      t = means[i];
      means[i] = means[191 - i];
      means[191 - i] = t;

      r = rms[i];
      rms[i] = rms[191 - i];
      rms[191 - i] = r;
      n++;
    }
  }
}

void unpack_process::threshold_to_gbt(vector<double> threshold)
{
    vector<double> th = threshold;

    for (int i = 128; i < th.size() + (asic_num * 3); i += 128 + 3) {
        th.insert(th.begin() + i, 3, 0);
    }

    long DataMSB = 0x0101;
    long DataLSB = 0x0102;
    string enable = "0x01020001";

    string wait = "&1";

    ostringstream datamsb;
    ostringstream datalsb;

    string bin_str;
    string bin_res;
    string bin_addr;
    string m_str;
    string l_str;
    string datamsb_str;
    string datalsb_str;

    string filename = file_write;
    ofstream output_file(filename, ios::out | ios::trunc);

    for (int i = 0; i < asic_num; i++)
    {

        for (int j = 0; j < 131; j++)
        {

            for (int k = 0; k < 2; k++)
            {

                string wr_msb_en = "0x0101400" + to_string(k);

                output_file << wr_msb_en << endl;
                output_file << enable << endl;

                bitset<24> bits((int)(th[j + i * 131]));
                bitset<8> bits3(j + 1);
                bin_str = bits.to_string();
                bin_addr = bits3.to_string();
                bin_res = bin_addr + bin_str;

                m_str = bin_res.substr(0, 16);
                l_str = bin_res.substr(16, 16);

                bitset<16> bits1(m_str);
                int msb = bits1.to_ulong();

                bitset<16> bits2(l_str);
                int lsb = bits2.to_ulong();

                datamsb << "0x" << setfill('0') << setw(4) << hex << DataMSB << setw(4) << hex << msb;
                datamsb_str = datamsb.str();

                datalsb << "0x" << setfill('0') << setw(4) << hex << DataMSB << setw(4) << hex << lsb;
                datalsb_str = datalsb.str();

                if (k == 0)
                {
                    output_file << datamsb_str << endl;
                }
                else
                {
                    output_file << datalsb_str << endl;
                }

                output_file << enable << endl;
                output_file << wait << endl;
                datamsb.str("");
                datalsb.str("");
            }
            output_file << "0x01014012" << endl;
            output_file << enable << endl;
            output_file << "0x01010000" << endl;
            output_file << enable << endl;
            output_file << wait << endl;

            // WREN EN
            output_file << "0x01015000" << endl;
            output_file << enable << endl;

            int en = 1;
            int wren = 1 << i;
            wren = wren & 0xffff;
            stringstream ss;
            ss << setw(4) << setfill('0') << hex << wren;
            string idx = ss.str();
            string ram_wren = "0x0101" + idx;
            output_file << ram_wren << endl;

            output_file << enable << endl;
            if (!(i == asic_num - 1 && j == 130))
            {
                output_file << wait << endl;
            }
        }
    }
    output_file.close();

    cout << cYELLOW << "threshold txt sucess output " << filename << endl;
    cout << endl;
}

/**
 * means: adc counts
 * rms: aC
 */
void unpack_process::draw(vector<vector<double>> means, vector<vector<double>> rms)
{
    string filename = "base";
    gStyle->SetImageScaling(3.);
    TCanvas *c1 = new TCanvas(filename.c_str(), filename.c_str());
    // 将画布分割成1行2列
    c1->Divide(2, 3);

    for (int i = 0; i < asic_num; i++) {
        TGraph *gr = new TGraph(means[i].size());
        for (int j = 0; j < means[i].size(); j++)
        {
            gr->SetPoint(j, j, means[i][j]);
        }

        c1->cd(i + 1); // 切换到第一个子画布
        gPad->SetGrid();
        string title = "Chip" + to_string(i);
        gr->SetTitle(title.c_str());
        gr->SetMarkerColor(4);
        gr->SetMarkerStyle(21);
        gr->SetMarkerSize(0.5);
        gr->Draw("APL");
        gr->GetXaxis()->SetRangeUser(0, 127);
    }
    string save_filename = save_path + "base.png";
    c1->SaveAs(save_filename.c_str());

    // rms
    string filename1 = "rms";
    gStyle->SetImageScaling(3.);
    TCanvas *c2 = new TCanvas(filename1.c_str(), filename1.c_str());
    c2->Divide(2, 3);

    for (int i = 0; i < asic_num; i++)
    {
        TGraph *gr = new TGraph(rms[i].size());
        for (int j = 0; j < rms[i].size(); j++)
        {
            gr->SetPoint(j, j, rms[i][j]);
        }

        c2->cd(i + 1); // 切换到第一个子画布
        gPad->SetGrid();
        string title = "Chip" + to_string(i);
        gr->SetTitle(title.c_str());
        gr->SetMarkerColor(4);
        gr->SetMarkerStyle(21);
        gr->SetMarkerSize(0.5);
        gr->Draw("APL");
        gr->GetXaxis()->SetRangeUser(0, 127);
    }
    string save_filename1 = save_path + "rms.png";
    c2->SaveAs(save_filename1.c_str());
}

/**
 * column: chip-wise
 * row: channel-wise (128)
 *
 * means: adc counts
 * rms: fC
 * threshold: adc counts (3*rms_counts)
 */
void unpack_process::write_base_txt(vector<vector<double>> means, vector<vector<double>> rms, vector<vector<double>> threshold)
{
    string filename = save_path + "base.txt";
    ofstream file(filename, ios::out | ios::trunc);

    for (int i = 0; i < means[0].size(); i++) {
        for (int j = 0; j < means.size(); j++) {
            file << means[j][i] << " ";
        }
        file << endl;
    }
    file.close();

    string filename1 = save_path + "rms.txt";
    ofstream file1(filename1, ios::out | ios::trunc);
    for (int i = 0; i < rms[0].size(); i++) {
        for (int j = 0; j < rms.size(); j++) {
            file1 << rms[j][i] << " ";
        }
        file1 << endl;
    }

    string filename2 = save_path + "thresh.txt";
    ofstream file2(filename2, ios::out | ios::trunc);
    for (int i = 0; i < threshold[0].size(); i++) {
        for (int j = 0; j < threshold.size(); j++) {
            file2 << threshold[j][i] << " ";
        }
        file2 << endl;
    }
}
